package com.bs.vectordatabase.controller;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.*;
import java.util.stream.Collectors;


@RestController
@RequiredArgsConstructor
@RequestMapping("/embedding/")
public class EmbeddingController {

    @Autowired
    private VectorStore vectorStore;

    /**
     * 系统提示词
     */
    private final static String SYSTEM_PROMPT = """
            你需要使用文档内容对用户提出的问题进行回复，同时你需要表现得天生就知道这些内容，
            不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要。 当用户提出的问题无法根据文档内容进行回复或者你也不清楚时，回复不知道即可。
            不能在回复中体现出你是根据给出的文档内容进行回复的，这点非常重要!
            文档内容如下:
            {documents}
                        
            """;


    @PostMapping("/add")
    public void add(@RequestBody CountUser countUser) {
        addContent(countUser);
    }

    @PostMapping("/update")
    public void update(@RequestBody CountUser countUser) {
        // 先删除
        vectorStore.delete(countUser.getAiKnowIds());

        // 插入
        addContent(countUser);
    }

    /**
     * 删除
     */
    @PostMapping("/remove")
    public void remove(@RequestBody CountUser countUser) {
        vectorStore.delete(countUser.getAiKnowIds());
    }

    @PostMapping("/query")
    public String query(@RequestBody CountUser countUser) {
        String message = countUser.getMessage();
        Long userId = countUser.getUserId();

        // 用户查询
        Filter.Expression filterExpression = new Filter.Expression(
                Filter.ExpressionType.EQ,
                new Filter.Key("userId"),
                new Filter.Value(userId)
        );
        // 问题查询
        SearchRequest request = SearchRequest.query(message)
                .withSimilarityThreshold(0.6)
                .withFilterExpression(filterExpression);

        List<Document> listOfSimilarDocuments = this.vectorStore.similaritySearch(request);
        if (listOfSimilarDocuments.size() == 0) {
            return "没有命中";
        } else {
            // 将Document列表中每个元素的content内容进行拼接获得documents
            String documentString = listOfSimilarDocuments.stream().map(Document::getContent).collect(Collectors.joining());
            Message systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentString));
            // 构建UserMessage对象
            UserMessage userMessage = new UserMessage(message);
            Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
            return prompt.getContents();
        }
    }


    private void addContent(CountUser countUser) {
        System.out.println("正在向量化数据");

        long start = System.currentTimeMillis(); // 开始时间
        String content = countUser.getContent();
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("userId", countUser.getUserId());
        List<Document> documents = List.of(
                new Document(countUser.getAiKnowId(), content, metadata));
        vectorStore.add(documents);

        long end = System.currentTimeMillis(); // 结束时间
        double elapsedSeconds = (end - start) / 1000.0; // 换算成秒
        System.out.println("向量化耗时：" + elapsedSeconds + " 秒");
    }



    /*@GetMapping("/embedding")
    public void embed(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));

        // Add the documents to PGVector
        vectorStore.add(documents);

        List<Document> listOfSimilarDocuments = this.vectorStore.similaritySearch(SearchRequest.query("Spring").withTopK(5));
        // 将Document列表中每个元素的content内容进行拼接获得documents
        String documentString = listOfSimilarDocuments.stream().map(Document::getContent).collect(Collectors.joining());
        Message systemMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentString));
        // 构建UserMessage对象
        UserMessage userMessage = new UserMessage(message);
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
        System.out.println(prompt);
    }*/
}
